#!/usr/bin/env python3
import pyrealsense2 as rs
import numpy as np
import cv2
from ultralytics import YOLO

# 加载YOLO模型（请确保模型路径正确）
# model = YOLO("yolo11n-seg.pt")
model = YOLO("we/sewer-yoloact-s.pt")

def main():
    # 创建 RealSense 管线并配置彩色流（640×480，30FPS）
    pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

    # 启动摄像头
    pipeline.start(config)
    print("RealSense 摄像头已启动，按 'q' 键退出。")
    
    # 预定义一组颜色，用于区分不同实例
    color_palette = [(0, 255, 0), (0, 0, 255), (255, 0, 0), 
                     (255, 255, 0), (0, 255, 255), (255, 0, 255)]
    
    try:
        while True:
            # 获取一帧数据
            frames = pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            if not color_frame:
                continue

            # 将 RealSense 帧转换为 numpy 数组
            color_image = np.asanyarray(color_frame.get_data())
            
            # 创建一个与原图尺寸相同的空白遮罩，用于累积各检测实例的分割掩码
            mask_overlay = np.zeros_like(color_image, dtype=np.uint8)
            # 用于存储边界框、类别与置信度信息，待后续在叠加后图像上绘制
            boxes_info = []
            
            # 调用YOLO模型进行实例分割预测，stream=True返回一个生成器
            results = model(color_image, stream=True)
            
            # 遍历预测结果（一般只返回一帧的检测结果）
            for r in results:
                if r.boxes is None:
                    continue
                # 当存在分割掩码时，从 r.masks.data 中提取（ultralytics中mask为张量）
                masks = r.masks.data if (r.masks is not None and hasattr(r.masks, "data")) else None
                # 对每个检测结果进行处理
                for i, box in enumerate(r.boxes):
                    # 取得边界框坐标（转换为整数）
                    xyxy = box.xyxy.cpu().numpy().astype(int)[0]
                    conf = box.conf.cpu().numpy()[0]
                    cls_id = int(box.cls.cpu().numpy()[0])
                    # 生成类别和置信度的文字标签
                    label = f"{model.names[cls_id]} {conf:.2f}"
                    # 根据检测序号选取颜色
                    color = color_palette[i % len(color_palette)]
                    boxes_info.append((xyxy, label, color))
                    
                    # 如果存在对应的掩码，则处理
                    if masks is not None:
                        # 取出第 i 个分割掩码，并转为 numpy 数组
                        mask = masks[i].cpu().numpy()
                        # 如果掩码存在多余的通道维度，则 squeeze
                        if mask.ndim == 3 and mask.shape[0] == 1:
                            mask = mask[0]
                        # 二值化：阈值设定为0.5
                        binary_mask = mask > 0.5
                        # 在 mask_overlay 中，将对应的区域填上指定颜色
                        mask_overlay[binary_mask] = color

            # 将原图与mask_overlay按透明度混合（alpha 值为0.5，可根据需要调整）
            annotated_image = cv2.addWeighted(color_image, 1.0, mask_overlay, 0.5, 0)
            
            # 在混合后的图像上绘制边界框和文字注释
            for (xyxy, label, color) in boxes_info:
                cv2.rectangle(annotated_image, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), color, 2)
                cv2.putText(annotated_image, label, (xyxy[0], xyxy[1] - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            
            # 显示标注后的图像
            cv2.imshow("RealSense Color Stream", annotated_image)

            # 按 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    finally:
        # 关闭摄像头和窗口
        pipeline.stop()
        cv2.destroyAllWindows()

if __name__ == '__main__':
    main()
